import os
import sys
import math
import threading
import time
from typing import Optional

import rospy
from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu


# 确保可导入同目录下的 test2.D435DetectModel
_CURRENT_DIR = os.path.dirname(__file__)
if _CURRENT_DIR not in sys.path:
    sys.path.append(_CURRENT_DIR)

from test2 import D435DetectModel  # noqa: E402


class KickballROSNode(D435DetectModel):
    """
    结合 ROS 的踢球跟踪节点：
    - IMU 订阅: /imu/data
    - 速度发布: /cmd_vel
    - 相机: 仍通过 pyrealsense2（由父类 D435DetectModel 持有并统一接口）

    通过继承 D435DetectModel 实现解耦：本类仅覆写速度下发与转向逻辑为 ROS 发布。
    """

    def __init__(
        self,
        model_path: str,
        target_label: str,
        prob_threshold: float = 0.8,
        cache_size: int = 10,
        timeout: float = 2.0,
        width: int = 640,
        height: int = 480,
        fps: int = 30,
        cmd_vel_topic: str = "/cmd_vel",
        imu_topic: str = "/imu/data",
        turn_yaw_rate: float = 0.6,  # rad/s
    ) -> None:
        super().__init__(
            yolo_path=model_path,
            prob_threshold=prob_threshold,
            cache_size=cache_size,
            timeout=timeout,
            width=width,
            height=height,
            fps=fps,
        )

        self.target_label = target_label
        self.cmd_vel_pub = rospy.Publisher(cmd_vel_topic, Twist, queue_size=10)
        self.imu_sub = rospy.Subscriber(imu_topic, Imu, self._imu_callback, queue_size=10)
        self._last_imu_msg = None  # type: Optional[Imu]
        self._turn_yaw_rate = abs(float(turn_yaw_rate))

        self._worker = None  # type: Optional[threading.Thread]
        self._shutdown = False

    # ============ ROS 相关 ============
    def _imu_callback(self, msg: Imu) -> None:
        self._last_imu_msg = msg

    def _publish_cmd_vel(self, vx: float, vy: float, yaw: float) -> None:
        tw = Twist()
        tw.linear.x = float(vx)
        tw.linear.y = float(vy)
        tw.angular.z = float(yaw)
        self.cmd_vel_pub.publish(tw)

    # ============ 抽象接口覆写 ============
    def send_velocity_command(self, vel_x: float, vel_y: float, yaw_z: float):
        # 应用父类限幅（若存在）
        try:
            if hasattr(self, "_limits"):
                vx = max(-self._limits['vx'], min(self._limits['vx'], vel_x))
                vy = max(-self._limits['vy'], min(self._limits['vy'], vel_y))
                yaw = max(-self._limits['yaw'], min(self._limits['yaw'], yaw_z))
            else:
                vx, vy, yaw = vel_x, vel_y, yaw_z
        except Exception:
            vx, vy, yaw = vel_x, vel_y, yaw_z

        self._publish_cmd_vel(vx, vy, yaw)

    def turn(self, angle_deg: float):
        # 简化的原地转向：以固定角速度旋转至目标角度
        yaw_rate = self._turn_yaw_rate if angle_deg >= 0 else -self._turn_yaw_rate
        duration = abs(math.radians(angle_deg)) / max(1e-6, self._turn_yaw_rate)

        end_time = time.time() + duration
        rate = rospy.Rate(50)
        while not rospy.is_shutdown() and time.time() < end_time and not self._shutdown:
            self._publish_cmd_vel(0.0, 0.0, yaw_rate)
            rate.sleep()
        # 停止转向
        self._publish_cmd_vel(0.0, 0.0, 0.0)

    # 可选：使用 IMU 改进父类 PID（当前父类未使用 IMU，此处仅传递存储的 IMU）
    def computeVelocity(self, rel_pos, imu_data):
        return super().computeVelocity(rel_pos, self._last_imu_msg)

    # ============ 运行与退出 ============
    def start(self):
        # 在独立线程中运行父类的持续跟踪管理逻辑
        self._worker = threading.Thread(target=self.GlobalManager, args=(self.target_label,), daemon=True)
        self._worker.start()

    def shutdown(self):
        self._shutdown = True
        try:
            # 停车
            self._publish_cmd_vel(0.0, 0.0, 0.0)
        except Exception:
            pass
        try:
            # 释放相机
            if hasattr(self, "camera") and self.camera is not None:
                self.camera.stop()
        except Exception:
            pass


def main():
    rospy.init_node("kickball_ros")

    # 读取参数
    model_path = rospy.get_param("~model_path", "yolov8n.pt")
    target_label = rospy.get_param("~target_label", "ball")
    prob_threshold = float(rospy.get_param("~prob_threshold", 0.8))
    cache_size = int(rospy.get_param("~cache_size", 10))
    timeout = float(rospy.get_param("~timeout", 2.0))
    width = int(rospy.get_param("~width", 640))
    height = int(rospy.get_param("~height", 480))
    fps = int(rospy.get_param("~fps", 30))
    cmd_vel_topic = rospy.get_param("~cmd_vel_topic", "/cmd_vel")
    imu_topic = rospy.get_param("~imu_topic", "/imu/data")
    turn_yaw_rate = float(rospy.get_param("~turn_yaw_rate", 0.6))

    node = KickballROSNode(
        model_path=model_path,
        target_label=target_label,
        prob_threshold=prob_threshold,
        cache_size=cache_size,
        timeout=timeout,
        width=width,
        height=height,
        fps=fps,
        cmd_vel_topic=cmd_vel_topic,
        imu_topic=imu_topic,
        turn_yaw_rate=turn_yaw_rate,
    )

    rospy.on_shutdown(node.shutdown)
    node.start()
    rospy.spin()


if __name__ == "__main__":
    main()


